import numpy as np
import SimpleITK as sitk
import glob,gc,os,sys,random
# ! <<< Outdated
# from keras.utils import to_categorical
import tensorflow as tf
from tensorflow.keras.utils import to_categorical
# ! >>>
from scipy.spatial.distance import cdist

sys.dont_write_bytecode = True

# ! <<< Self modify for relative import by Eval process
# import dataBase as DB
# from paraClass import *
if len(__name__.split('.')) > 1:
    sys.path.append('.')
    from . import dataBase as DB
    from .paraClass import *
else:
    import dataBase as DB
    from paraClass import *
# ! >>>

def read_txt_line(path):
    with open(path) as fid:
        lines = fid.readlines()
    return lines

def label_adjustment(label, nclass, negids = None, sortids = None):
    if negids is None:
        negids = []
    if sortids is None:
        sortids = []
    for i in negids:
        label[label==i] = 0
    if nclass==2:
        label[label!=0] = 1
    if len(sortids)!=0:
        labelcopy = label.copy()
        for tab, i in enumerate(sortids, start=1):
            label[labelcopy==i] = tab
    return label

def read_file_under_folder(path,suffix,txtpath=None):
    return (
        glob.glob(f'{path}/*.{suffix}')
        if txtpath is None
        else read_txt_line(txtpath)
    )

def sample_random_patch_2D(img,patchdims,samplenum,mask):
    leng,col,row = img.shape

    patches = np.empty(shape=(0,patchdims[0],patchdims[1],patchdims[2]), dtype=np.float32)
    patchcoors = np.zeros((1,6),dtype=np.uint16)

    nostart = 0
    zstep = int(patchdims[0]/2)

    for i in range(samplenum):
        zs = np.random.randint(0,leng)
        ys = np.random.randint(0, np.maximum(col - patchdims[1], 1))
        xs = np.random.randint(0, np.maximum(row - patchdims[2], 1))

        ze = zs + patchdims[0]
        ye = ys + patchdims[1]
        xe = xs + patchdims[2]

        curpatch = img[zs:ze,ys:ye,xs:xe][:].reshape([1,patchdims[0],patchdims[1],patchdims[2]])
        patches = np.append(patches,curpatch,axis=0)

        if nostart == 0:
            patchcoors[0,:] = [zs,ze,ys,ye,xs,xe]
            nostart = 1
        else:
            patchcoors = np.row_stack((patchcoors,[zs,ze,ys,ye,xs,xe]))

    patches = patches.astype('float32')
    return patches,patchcoors

def sample_random_patch_3D(img,patchdims,samplenum,mask=None):
    leng,col,row = img.shape
    patches = np.empty(shape=(0,1,patchdims[0],patchdims[1],patchdims[2]), dtype=np.float32)
    patchcoors = np.zeros((1,6),dtype=np.uint16)

    nostart = 0
    counter = 0
    while counter < samplenum:
        zs = np.random.randint(0,leng-patchdims[0])
        ys = np.random.randint(0,np.maximum(col-patchdims[1],1))
        xs = np.random.randint(0,np.maximum(row-patchdims[2],1))

        ze = zs + patchdims[0]
        ye = ys + patchdims[1]
        xe = xs + patchdims[2]

        if mask is not None:
            curpatchmask = mask[zs:ze,ys:ye,xs:xe]
            if np.count_nonzero(curpatchmask)==0:
                continue

        curpatch = img[zs:ze,ys:ye,xs:xe][:].reshape([1,1,patchdims[0],patchdims[1],patchdims[2]])
        patches = np.append(patches,curpatch,axis=0)
        if nostart == 0:
            patchcoors[0,:] = [zs,ze,ys,ye,xs,xe]
            nostart = 1
        else:
            patchcoors = np.row_stack((patchcoors,[zs,ze,ys,ye,xs,xe]))
        counter +=1

    patches = patches.astype(np.float32)
    return patches, patchcoors

def online_batch_generator_2D(imgpath,labelpath,batchnum,sampleparas,normparas,txtpath=None,axis=-1):
    imglist = read_file_under_folder(imgpath,'nii.gz',txtpath)
    print(f'The total number of images is: {len(imglist)}')

    while 1:
        ## shuffle
        accumtab = 0
        random.shuffle(imglist)

        imgs_train = np.empty(shape=(0,sampleparas.patchdims[0],sampleparas.patchdims[1],sampleparas.patchdims[2]),dtype=np.float32)
        # imgs_label_train = np.empty(shape=(0,sampleparas.patchlabeldims[0],sampleparas.patchlabeldims[1],sampleparas.patchlabeldims[2]),dtype=np.uint8)
        imgs_label_train = np.empty(shape=(0,sampleparas.patchlabeldims[0],sampleparas.patchlabeldims[1],sampleparas.patchlabeldims[2]),dtype=np.float32)

        for imgfile in imglist:
            imgpath = imgfile.split('|')[0]
            mskpath = imgfile.split('|')[1].split('\n')[0]
            # midname = imgfile[imgfile.rindex("/")+1:]
            # if midname.count(".")>0:
            #     midname = midname[0:midname.index('.')]
            # else:
            #     midname = midname.strip('\n')
            # gc.collect()

            # imgraw = sitk.ReadImage(imgpath + "/" + midname + ".nii.gz")
            imgraw = sitk.ReadImage(imgpath)
            img = sitk.GetArrayFromImage(imgraw)
            # if not os.path.exists(labelpath + "/" + midname + "_MASK.nii.gz"):
            #     continue
            # labelraw = sitk.ReadImage(labelpath + "/" + midname + "_MASK.nii.gz")
            labelraw = sitk.ReadImage(mskpath)
            label = sitk.GetArrayFromImage(labelraw)
            if len(np.unique(label))==1:
                continue
            label = label_adjustment(label,sampleparas.nclass,sampleparas.negids,sampleparas.sortids)

            if img.shape!=label.shape:
                continue

            # normalization of the input images
            normimg = DB.min_max_normalization(img,normparas.lmin, normparas.rmax)
            patches,patchcoors = sample_random_patch_2D(normimg,sampleparas.patchdims,sampleparas.samplenum,label)
            patchlabeles = DB.sample_label_with_coor_2D(label,sampleparas.patchlabeldims,patchcoors,sampleparas.nclass)

            imgs_train = np.append(imgs_train,patches,axis=0)
            imgs_label_train = np.append(imgs_label_train,patchlabeles,axis=0)

            accumtab += 1
            if accumtab == batchnum:
                gc.collect()
                accumtab = 0

                # normalization for the whole sample data
                if sampleparas.meanvalue is not None:
                    imgs_train = imgs_train - sampleparas.meanvalue

                if axis == -1:
                    imgs_train = np.transpose(imgs_train, (0,2,3,1))
                    imgs_label_train = np.transpose(imgs_label_train, (0,2,3,1))

                yield (imgs_train, imgs_label_train)
                imgs_train = np.empty(
                    shape=(0, sampleparas.patchdims[0], sampleparas.patchdims[1], sampleparas.patchdims[2]),
                    dtype=np.float32)
                imgs_label_train = np.empty(shape=(
                0, sampleparas.patchlabeldims[0], sampleparas.patchlabeldims[1], sampleparas.patchlabeldims[2]),
                                            dtype=np.uint8)

def online_batch_generator_3D(imgpath,labelpath,batchnum,sampleparas,normparas,txtpath=None,axis=-1):
    # ! <<<
    # imglist = iof.read_file_under_folder(imgpath,'nii.gz',txtpath)
    imglist = read_file_under_folder(imgpath,'nii.gz',txtpath)
    # ! >>>
    print("The total number of images in %s is: %d"%(imgpath,len(imglist)))

    while 1:
        ## shuffle
        accumtab = 0
        random.shuffle(imglist)

        imgs_train = np.empty(shape=(0,1,sampleparas.patchdims[0],sampleparas.patchdims[1],sampleparas.patchdims[2]),dtype=np.float32)
        imgs_label_train = np.empty(shape=(0,1,sampleparas.patchlabeldims[0],sampleparas.patchlabeldims[1],sampleparas.patchlabeldims[2]),dtype=np.uint8)

        for imgfile in imglist:
            midname = imgfile[imgfile.rindex("/")+1:]
            # remove the line break character
            if midname.count(".")>0:
                midname = midname[0:midname.index('.')]
            else:
                midname = midname.strip('\n')
            gc.collect()

            imgraw = sitk.ReadImage(imgpath + "/" + midname + ".nii.gz")
            img = sitk.GetArrayFromImage(imgraw)

            if not os.path.exists(labelpath + "/" + midname + "_MASK.nii.gz"):
                continue
            labelraw = sitk.ReadImage(labelpath + "/" + midname + "_MASK.nii.gz")
            label = sitk.GetArrayFromImage(labelraw)
            if len(np.unique(label))==1:
                continue
            # ! <<<
            # label = ltlor.label_adjustment(label,sampleparas.nclass,sampleparas.negids,sampleparas.sortids)
            label = label_adjustment(label,sampleparas.nclass,sampleparas.negids,sampleparas.sortids)
            # ! >>>

            if img.shape!=label.shape:
                continue

            # normalization of the input images
            normimg = DB.min_max_normalization(img,normparas.lmin, normparas.rmax)
            patches,patchcoors = sample_random_patch_3D(normimg,sampleparas.patchdims,sampleparas.samplenum,label)
            # ! <<<
            # patchlabeles = DB.sample_label_with_coor_3D(label,sampleparas.patchlabeldims,patchcoors,sampleparas.nclass)
            patchlabeles = DB.sample_label_with_coor_2D(label,sampleparas.patchlabeldims,patchcoors,sampleparas.nclass, model_dimension=3)
            # ! >>>

            imgs_train = np.append(imgs_train,patches,axis=0)
            imgs_label_train = np.append(imgs_label_train, patchlabeles,axis=0)

            accumtab += 1
            if accumtab == batchnum:
                gc.collect()
                accumtab = 0

                # normalization for the whole sample data
                if sampleparas.meanvalue is not None:
                    imgs_train = imgs_train - sampleparas.meanvalue

                if axis == -1:
                    imgs_train = np.transpose(imgs_train, (0,2,3,4,1))
                    imgs_label_train = np.transpose(imgs_label_train, (0,2,3,4,1))

                yield (imgs_train, imgs_label_train)
                imgs_train = np.empty(
                    shape=(0, 1, sampleparas.patchdims[0], sampleparas.patchdims[1], sampleparas.patchdims[2]),
                    dtype=np.float32)
                imgs_label_train = np.empty(shape=(
                0, 1, sampleparas.patchlabeldims[0], sampleparas.patchlabeldims[1], sampleparas.patchlabeldims[2]),
                                            dtype=np.uint8)
